import numpy as np
from collections import Counter

class KNearestNeighbors:
    def __init__(self, k=3, distance_metric='euclidean'):
        
        self.k = k
        self.distance_metric = distance_metric
        self.X_train = None
        self.y_train = None
    
    def fit(self, X, y):
       
        self.X_train = np.array(X)
        self.y_train = np.array(y)
        return self
    
    def _calculate_distance(self, x1, x2):
       
        if self.distance_metric == 'euclidean':
            # 欧几里得距离
            return np.sqrt(np.sum((x1 - x2) ** 2))
        elif self.distance_metric == 'manhattan':
            # 曼哈顿距离
            return np.sum(np.abs(x1 - x2))
        elif self.distance_metric == 'minkowski':
            # 闵可夫斯基距离 (p=3)
            return np.sum(np.abs(x1 - x2) ** 3) ** (1/3)
        else:
            raise ValueError("不支持的距離度量方法")
    
    def _get_neighbors(self, x):
      
        distances = []
        
        # 计算与所有训练样本的距离
        for i, train_sample in enumerate(self.X_train):
            dist = self._calculate_distance(x, train_sample)
            distances.append((i, dist))
        
        # 按距离排序并获取前k个邻居
        distances.sort(key=lambda x: x[1])
        neighbors = distances[:self.k]
        
        return neighbors
    
    def predict_single(self, x):
      
        neighbors = self._get_neighbors(x)
        
        # 获取邻居的标签
        neighbor_labels = [self.y_train[i] for i, _ in neighbors]
        
        # 返回最常见的标签
        most_common = Counter(neighbor_labels).most_common(1)
        return most_common[0][0]
    
    def predict(self, X):
      
        X = np.array(X)
        predictions = []
        
        for sample in X:
            pred = self.predict_single(sample)
            predictions.append(pred)
        
        return np.array(predictions)
    
    def score(self, X, y):
      
        predictions = self.predict(X)
        return np.mean(predictions == y)
    
    def predict_proba(self, X):
     
        X = np.array(X)
        proba = []
        
        for sample in X:
            neighbors = self._get_neighbors(sample)
            neighbor_labels = [self.y_train[i] for i, _ in neighbors]
            
            # 计算每个类别的比例作为概率
            label_counts = Counter(neighbor_labels)
            unique_labels = sorted(label_counts.keys())
            sample_proba = [label_counts.get(label, 0) / self.k for label in unique_labels]
            proba.append(sample_proba)
        
        return np.array(proba)